import pybullet as p
import pybullet_data
import numpy as np


class PybulletPlanner:

    def __init__(self, scene, urdf, move_group='tool0'):
        self._scene = scene
        self.physicsClient = p.connect(p.DIRECT)
        
        # 加载 URDF 机械臂模型
        self.ROBOT_ID = p.loadURDF(urdf, useFixedBase=True, physicsClientId=self.physicsClient)
        self.TOOL0_LINK_INDEX = 12

        self.num_joints = p.getNumJoints(self.ROBOT_ID, physicsClientId=self.physicsClient)
        self.joint_info = [(p.getJointInfo(self.ROBOT_ID, i), i) for i in range(self.num_joints)]
        self.movable_joints = [info[1] for info in self.joint_info if p.getJointInfo(self.ROBOT_ID, info[1])[2] != p.JOINT_FIXED]

    def solve_fk_homogeneous(self, qpos):
        '''
        Input qpos
        return the target_to_base_homogenious_matrix
        '''
        joint_positions = qpos
        for movable_joint_id, joint_position in zip(self.movable_joints, joint_positions):
            p.resetJointState(self.ROBOT_ID, movable_joint_id, joint_position, physicsClientId=self.physicsClient)
        # 获取机械臂ee link状态
        link_state = p.getLinkState(self.ROBOT_ID, self.TOOL0_LINK_INDEX, physicsClientId=self.physicsClient)
        xyz_position = link_state[0]  # ee link的世界坐标系中的位置
        orientation = link_state[1]  # ee link的世界坐标系中的四元数姿态
        T = self.create_homogeneous_matrix(xyz_position, orientation)
        return T 
    
    def create_homogeneous_matrix(self, position, quaternion):
        # 创建4x4的齐次变换矩阵
        R = self.quaternion_to_rotation_matrix(quaternion)  # 旋转矩阵
        T = np.eye(4)  # 创建4x4单位矩阵
        T[:3, :3] = R  # 将旋转矩阵放入T的上左部分
        T[:3, 3] = position  # 将平移部分放入T的右上部分
        return T
    
    def quaternion_to_rotation_matrix(self, quaternion):
        '''
        Input quaternion [x y z w]
        Output rotation matrix R
        '''
        # 四元数到旋转矩阵的转换
        x,y,z,w= quaternion
        R = np.array([[1 - 2 * (y**2 + z**2), 2 * (x * y - z * w), 2 * (x * z + y * w)],
                    [2 * (x * y + z * w), 1 - 2 * (x**2 + z**2), 2 * (y * z - x * w)],
                    [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x**2 + y**2)]])
        return R
    
    def solve_fk_velocity(self, joint_rad, joint_vel):
        """
        计算末端执行器在工具坐标系下的速度
        :param joint_rad: 关节角度 (弧度)
        :param joint_vel: 关节角速度
        :return: 末端执行器速度 (线速度 + 角速度) 在工具坐标系 (6x1)
        """
        # 计算正向运动学变换矩阵
        T = self.solve_fk_homogeneous(joint_rad)
        R = T[:3, :3]  # 旋转矩阵

        # 计算EE到 Base 变换矩阵
        T_EE_to_Base = np.eye(6)
        T_EE_to_Base[:3, :3] = R
        T_EE_to_Base[3:, 3:] = R

        # 计算 EE 到基座的逆变换
        T_Base_to_EE = np.linalg.inv(T_EE_to_Base)
        num_movable_joints = len(self.movable_joints)
        # 计算雅可比矩阵
        joint_accelerations = [0] * num_movable_joints # 关节加速度设为 0
        joint_velocity = [0] * num_movable_joints
        assert len(joint_rad) == num_movable_joints, f"joint_rad 长度错误，期望 {num_movable_joints}，但收到 {len(joint_rad)}"
        
        J_linear, J_angular = p.calculateJacobian(self.ROBOT_ID, self.TOOL0_LINK_INDEX, [0, 0, 0], 
                                                joint_rad.tolist(), joint_velocity, joint_accelerations, physicsClientId=self.physicsClient)

        # 组合线速度和角速度雅可比矩阵
        J = np.vstack((np.array(J_linear), np.array(J_angular)))

        # 计算基座坐标系下的末端速度
        V_base_Coord = J @ np.array(joint_vel).reshape(-1,1)  # 确保维度匹配

        # 变换到工具坐标系
        # V_EE_Coord = T_Base_to_EE @ V_base_Coord

        return V_base_Coord


if __name__ == "__main__":
    # Set up the PyBullet physics client.
    # You can use p.GUI to visualize the robot or p.DIRECT for headless computation.
    physicsClient = p.connect(p.DIRECT)

    # Add search path for common URDF files (like plane.urdf).
    p.setAdditionalSearchPath(pybullet_data.getDataPath())

    # Load your robot's URDF file. Replace 'kuka_lbr.urdf' with your file.
    robot_id = p.loadURDF("kuka_lbr.urdf", useFixedBase=True)

    # Define the joint positions (qpos) and velocities (qvel)
    # These are the values you would get from your control system or simulation.
    # Make sure the number of values matches the number of joints in your URDF.
    num_joints = p.getNumJoints(robot_id)
    qpos = np.zeros(num_joints)  # Replace with your actual joint positions
    qvel = np.zeros(num_joints)  # Replace with your actual joint velocities

    # Set the joint states (position and velocity) of the robot
    for i in range(num_joints):
        p.resetJointState(robot_id, i, qpos[i], qvel[i])

    # Specify the index of the link you want to get the state for.
    # You can use getJointInfo to find the link index. For example, to get
    # the index of the link named 'lbr_iiwa_link_7':
    # for j in range(num_joints):
    #     info = p.getJointInfo(robot_id, j)
    #     if info[12].decode('utf-8') == 'lbr_iiwa_link_7':
    #         link_index = info[0]
    #         break

    # For this example, let's assume the link_index is 6 (a common end-effector link index).
    link_index = 6

    # Get the link state, which includes position and velocity
    link_state = p.getLinkState(robot_id, link_index, computeLinkVelocity=1, computeForwardKinematics=1)

    # The result is a tuple. The first element is the position, and the seventh is the linear velocity.
    # See PyBullet documentation for full details:
    # https://docs.google.com/document/d/10sXEhzFRSnvNCjfTLEpRYehaIPyScCQf/edit#heading=h.g7d2062x317x
    link_position = link_state[0]
    link_linear_velocity = link_state[6]

    print(f"Position of link {link_index}: {link_position}")
    print(f"Linear velocity of link {link_index}: {link_linear_velocity}")

    # Disconnect from the physics client
    p.disconnect()